// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>
#include <cstring>
#include <memory>
#include <type_traits>

#include <nebula/array/builder_base.h>
#include <nebula/core/buffer.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <turbo/base/macros.h>

namespace nebula {

    /// \addtogroup numeric-builders
    ///
    /// @{

    namespace internal {

        class TURBO_EXPORT AdaptiveIntBuilderBase : public ArrayBuilder {
        public:
            AdaptiveIntBuilderBase(uint8_t start_int_size, MemoryPool *pool,
                                   int64_t alignment = kDefaultBufferAlignment);

            explicit AdaptiveIntBuilderBase(MemoryPool *pool,
                                            int64_t alignment = kDefaultBufferAlignment)
                    : AdaptiveIntBuilderBase(sizeof(uint8_t), pool, alignment) {}

            /// \brief append multiple nulls
            /// \param[in] length the number of nulls to append
            turbo::Status append_nulls(int64_t length) final {
                TURBO_RETURN_NOT_OK(CommitPendingData());
                if (TURBO_LIKELY(length > 0)) {
                    TURBO_RETURN_NOT_OK(Reserve(length));
                    memset(data_->mutable_data() + length_ * int_size_, 0, int_size_ * length);
                    unsafe_set_null(length);
                }
                return turbo::OkStatus();
            }

            turbo::Status append_null() final {
                pending_data_[pending_pos_] = 0;
                pending_valid_[pending_pos_] = 0;
                pending_has_nulls_ = true;
                ++pending_pos_;
                ++length_;
                ++null_count_;

                if (TURBO_UNLIKELY(pending_pos_ >= pending_size_)) {
                    return CommitPendingData();
                }
                return turbo::OkStatus();
            }

            turbo::Status append_empty_values(int64_t length) final {
                TURBO_RETURN_NOT_OK(CommitPendingData());
                if (TURBO_LIKELY(length > 0)) {
                    TURBO_RETURN_NOT_OK(Reserve(length));
                    memset(data_->mutable_data() + length_ * int_size_, 0, int_size_ * length);
                    unsafe_set_not_null(length);
                }
                return turbo::OkStatus();
            }

            turbo::Status append_empty_value() final {
                pending_data_[pending_pos_] = 0;
                pending_valid_[pending_pos_] = 1;
                ++pending_pos_;
                ++length_;

                if (TURBO_UNLIKELY(pending_pos_ >= pending_size_)) {
                    return CommitPendingData();
                }
                return turbo::OkStatus();
            }

            void reset() override;

            turbo::Status resize(int64_t capacity) override;

        protected:
            turbo::Status AppendInternal(const uint64_t val) {
                pending_data_[pending_pos_] = val;
                pending_valid_[pending_pos_] = 1;
                ++pending_pos_;
                ++length_;

                if (TURBO_UNLIKELY(pending_pos_ >= pending_size_)) {
                    return CommitPendingData();
                }
                return turbo::OkStatus();
            }

            virtual turbo::Status CommitPendingData() = 0;

            template<typename new_type, typename old_type>
            typename std::enable_if<sizeof(old_type) >= sizeof(new_type), turbo::Status>::type
            ExpandIntSizeInternal();

            template<typename new_type, typename old_type>
            typename std::enable_if<(sizeof(old_type) < sizeof(new_type)), turbo::Status>::type
            ExpandIntSizeInternal();

            std::shared_ptr<ResizableBuffer> data_;
            uint8_t *raw_data_ = nullptr;

            const uint8_t start_int_size_;
            uint8_t int_size_;

            static constexpr int32_t pending_size_ = 1024;
            uint8_t pending_valid_[pending_size_];
            uint64_t pending_data_[pending_size_];
            int32_t pending_pos_ = 0;
            bool pending_has_nulls_ = false;
        };

    }  // namespace internal

    class TURBO_EXPORT AdaptiveUIntBuilder : public internal::AdaptiveIntBuilderBase {
    public:
        explicit AdaptiveUIntBuilder(uint8_t start_int_size,
                                     MemoryPool *pool = default_memory_pool());

        explicit AdaptiveUIntBuilder(MemoryPool *pool = default_memory_pool())
                : AdaptiveUIntBuilder(sizeof(uint8_t), pool) {}

        using internal::AdaptiveIntBuilderBase::reset;

        /// Scalar append
        turbo::Status append(const uint64_t val) { return AppendInternal(val); }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a contiguous C array of values
        /// \param[in] length the number of values to append
        /// \param[in] valid_bytes an optional sequence of bytes where non-zero
        /// indicates a valid (non-null) value
        /// \return turbo::Status
        turbo::Status append_values(const uint64_t *values, int64_t length,
                                   const uint8_t *valid_bytes = nullptr);

        turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) override;

        std::shared_ptr<DataType> type() const override;

    protected:
        turbo::Status CommitPendingData() override;

        turbo::Status ExpandIntSize(uint8_t new_int_size);

        turbo::Status AppendValuesInternal(const uint64_t *values, int64_t length,
                                           const uint8_t *valid_bytes);

        template<typename new_type>
        turbo::Status ExpandIntSizeN();
    };

    class TURBO_EXPORT AdaptiveIntBuilder : public internal::AdaptiveIntBuilderBase {
    public:
        explicit AdaptiveIntBuilder(uint8_t start_int_size,
                                    MemoryPool *pool = default_memory_pool(),
                                    int64_t alignment = kDefaultBufferAlignment);

        explicit AdaptiveIntBuilder(MemoryPool *pool = default_memory_pool(),
                                    int64_t alignment = kDefaultBufferAlignment)
                : AdaptiveIntBuilder(sizeof(uint8_t), pool, alignment) {}

        using internal::AdaptiveIntBuilderBase::reset;

        /// Scalar append
        turbo::Status append(const int64_t val) { return AppendInternal(static_cast<uint64_t>(val)); }

        /// \brief append a sequence of elements in one shot
        /// \param[in] values a contiguous C array of values
        /// \param[in] length the number of values to append
        /// \param[in] valid_bytes an optional sequence of bytes where non-zero
        /// indicates a valid (non-null) value
        /// \return turbo::Status
        turbo::Status append_values(const int64_t *values, int64_t length,
                                   const uint8_t *valid_bytes = nullptr);

        turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) override;

        std::shared_ptr<DataType> type() const override;

    protected:
        turbo::Status CommitPendingData() override;

        turbo::Status ExpandIntSize(uint8_t new_int_size);

        turbo::Status AppendValuesInternal(const int64_t *values, int64_t length,
                                           const uint8_t *valid_bytes);

        template<typename new_type>
        turbo::Status ExpandIntSizeN();
    };

    /// @}

}  // namespace nebula
